import os
import time
import uuid
import torch.serialization
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
from pydantic import BaseModel
from typing import Optional
from infer.modules.vc.modules import VC
from configs.config import Config



# 补丁：强制设置 weights_only=False
_orig_load = torch.load
def patched_load(*args, **kwargs):
    kwargs['weights_only'] = False
    return _orig_load(*args, **kwargs)
torch.load = patched_load

class VoiceSession:
    def __init__(self, config):
        self.vc = VC(config)
        self.last_access = time.time()
        self.model_loaded = False
    
    def update_access_time(self):
        self.last_access = time.time()

class SessionManager:
    def __init__(self):
        self.sessions = {}
        self.config = Config()
        
    def create_session(self, session_id: str = None) -> str:
        if session_id is None:
            session_id = str(uuid.uuid4())
        if session_id not in self.sessions:
            self.sessions[session_id] = VoiceSession(self.config)
        return session_id
    
    def get_session(self, session_id: str) -> Optional[VoiceSession]:
        session = self.sessions.get(session_id)
        if session:
            session.update_access_time()
        return session
    
    def cleanup_sessions(self, timeout: int = 3600):
        current_time = time.time()
        for sid in list(self.sessions.keys()):
            if current_time - self.sessions[sid].last_access > timeout:
                del self.sessions[sid]

class VoiceRequest(BaseModel):
    session_id: Optional[str] = None
    model_path: str
    protect0: float = 0.33  # 添加保护参数0
    protect1: float = 0.33  # 添加保护参数1

    
class ConversionRequest(BaseModel):
    session_id: str
    input_path: str  # 输入音频文件的绝对路径
    f0_up_key: int = 0  # 变调参数,正数升调,负数降调,以半音为单位。例如:12表示升高一个八度,-12表示降低一个八度
    f0_method: str = "rmvpe"  # 音高提取算法,建议使用 rmvpe
    index_rate: float = 0.5  # 检索特征占比
    protect: float = 0.33  # 保护清辅音强度
    filter_radius: int = 3  # 滤波器半径
    resample_sr: int = 0  # 重采样率,设置为0表示不进行重采样
    rms_mix_rate: float = 0.25  # 音量包络替换比例
    wav_out_put: str = '/data1/gitee/retrieval-based-voice-conversion-web-ui/audio/out' #音频输出地址

app = FastAPI()
session_manager = SessionManager()


@app.post("/loadweight_session")
async def create_session(request: VoiceRequest):
    try:
        session_id = session_manager.create_session(request.session_id)
        session = session_manager.get_session(session_id)
        
        # 加载模型,传入protect参数
        result = session.vc.get_vc(
            request.model_path,
            request.protect0,
            request.protect1
        )
        session.model_loaded = True
        
        return {
            "status": "success",
            "session_id": session_id,
            "message": f"Model {request.model_path} loaded successfully",
            "result": result  # 返回get_vc的结果
        }
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }


@app.post("/convert_voice")
async def convert_voice(
    request: ConversionRequest
):
    try:
        session = session_manager.get_session(request.session_id)
        if not session:
            return {"status": "error", "message": "Session not found"}
        if not session.model_loaded:
            return {"status": "error", "message": "No model loaded for this session"}
            
        # 检查输入文件是否存在
        if not os.path.exists(request.input_path):
            return {"status": "error", "message": f"Input file not found: {request.input_path}"}
        
        try:
            # 执行语音转换
            info, audio = session.vc.vc_single(
                sid=0,  # 因为模型已经加载,这里sid可以是固定值
                input_audio_path=request.input_path,  # 直接使用输入文件路径
                f0_up_key=request.f0_up_key,
                f0_file=None,
                f0_method=request.f0_method,
                file_index="",
                file_index2="",
                index_rate=request.index_rate,
                filter_radius=request.filter_radius,
                resample_sr=request.resample_sr,
                rms_mix_rate=request.rms_mix_rate,
                protect=request.protect
            )
            
            if "Success" in info:
                # 在输出目录下创建会话专属目录
                output_dir = os.path.join(request.wav_out_put, request.session_id)
                os.makedirs(output_dir, exist_ok=True)
                
                output_filename = f"output_{os.path.basename(request.input_path)}"
                output_path = os.path.join(output_dir, output_filename)
                
                sr, audio_data = audio
                import soundfile as sf
                sf.write(output_path, audio_data, sr)
                
                return {
                    "status": "success",
                    "message": info,
                    "output_path": output_path
                }
            else:
                return {
                    "status": "error",
                    "message": info
                }
                
        except Exception as e:
            return {
                "status": "error",
                "message": str(e)
            }
                
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }

@app.get("/download_audio")
async def download_audio(file_path: str):
    """
    下载音频文件接口
    
    参数:
        file_path: 音频文件的完整路径
    返回:
        文件下载响应
    """
    try:
        if not os.path.exists(file_path):
            return {
                "status": "error",
                "message": f"文件不存在: {file_path}"
            }
            
        # 获取文件名
        file_name = os.path.basename(file_path)
        
        # 返回文件下载响应
        return FileResponse(
            path=file_path,
            filename=file_name,
            media_type="audio/wav"
        )
        
    except Exception as e:
        return {
            "status": "error",
            "message": str(e)
        }



# 定期清理过期会话的任务
@app.on_event("startup")
async def startup_event():
    import asyncio
    async def cleanup_task():
        while True:
            session_manager.cleanup_sessions()
            await asyncio.sleep(1800)  # 每30分钟清理一次
    asyncio.create_task(cleanup_task())

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7869)